package core

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/cogentcore/webgpu/wgpu"
)

// Context 管理WebGPU设备和资源
type Context struct {
	instance *wgpu.Instance
	adapter  *wgpu.Adapter
	device   *wgpu.Device
	queue    *wgpu.Queue
	pool     *MemoryPool
	limits   wgpu.Limits

	// 轮询守护协程
	pollQuit   chan struct{}  // 退出信号
	pollActive sync.WaitGroup // 确保干净退出
	pollOnce   sync.Once      // 确保只启动一次

	// 内存池和内核缓存
	kernelCache sync.Map
}

// ContextConfig 配置选项
type ContextConfig struct {
	PowerPreference   wgpu.PowerPreference
	EnableValidation  bool
	MemoryPoolSize    uint64
	DisablePollDaemon bool          // 禁用守护协程
	PollInterval      time.Duration // 轮询频率
}

// WithValidation 启用调试验证
func WithValidation(enabled bool) func(*ContextConfig) {
	return func(c *ContextConfig) {
		c.EnableValidation = enabled
	}
}

// WithPowerPreference 设置电源偏好
func WithPowerPreference(pref wgpu.PowerPreference) func(*ContextConfig) {
	return func(c *ContextConfig) {
		c.PowerPreference = pref
	}
}

// WithMemoryPoolSize 设置内存池大小
func WithMemoryPoolSize(size uint64) func(*ContextConfig) {
	return func(c *ContextConfig) {
		c.MemoryPoolSize = size
	}
}

// NewContext 创建GPU上下文（包含守护协程）
func NewContext(opts ...func(*ContextConfig)) (*Context, error) {
	config := &ContextConfig{
		PowerPreference:   wgpu.PowerPreferenceHighPerformance,
		EnableValidation:  false,
		MemoryPoolSize:    128 * 1024 * 1024,
		DisablePollDaemon: false,
		PollInterval:      100 * time.Microsecond,
	}

	for _, opt := range opts {
		opt(config)
	}

	instance := wgpu.CreateInstance(nil)
	if instance == nil {
		return nil, errors.New("failed to create WebGPU instance")
	}

	adapter, err := instance.RequestAdapter(&wgpu.RequestAdapterOptions{
		PowerPreference: config.PowerPreference,
	})
	if err != nil {
		instance.Release()
		return nil, fmt.Errorf("failed to request adapter: %v", err)
	}

	device, err := adapter.RequestDevice(&wgpu.DeviceDescriptor{
		Label:          "wgmat-device",
		RequiredLimits: nil,
	})
	if err != nil {
		adapter.Release()
		instance.Release()
		return nil, fmt.Errorf("failed to request device: %v", err)
	}

	queue := device.GetQueue()
	supportedLimits := device.GetLimits()
	pool := NewMemoryPool(device, config.MemoryPoolSize)

	ctx := &Context{
		instance: instance,
		adapter:  adapter,
		device:   device,
		queue:    queue,
		pool:     pool,
		limits:   supportedLimits.Limits,
		pollQuit: make(chan struct{}),
	}

	// 自动启动轮询守护协程
	if !config.DisablePollDaemon {
		ctx.pollOnce.Do(func() {
			ctx.pollActive.Add(1)
			go ctx.pollDaemon(config.PollInterval)
		})
	}

	return ctx, nil
}

// pollDaemon 轮询守护协程（驱动GPU回调）
func (c *Context) pollDaemon(interval time.Duration) {
	defer c.pollActive.Done()
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-c.pollQuit:
			return
		case <-ticker.C:
			c.device.Poll(false, nil)
		}
	}
}

// Device 返回WebGPU设备
func (c *Context) Device() *wgpu.Device {
	return c.device
}

// Queue 返回命令队列
func (c *Context) Queue() *wgpu.Queue {
	return c.queue
}

// MemoryPool 返回内存池
func (c *Context) MemoryPool() *MemoryPool {
	return c.pool
}

// KernelCache 返回内核缓存
func (c *Context) KernelCache() *sync.Map {
	return &c.kernelCache
}

// Close 释放资源（等待守护协程退出）
func (c *Context) Close() error {
	// 停止守护协程
	if c.pollQuit != nil {
		close(c.pollQuit)
		c.pollActive.Wait()
	}

	// 清理内核缓存
	c.kernelCache.Range(func(key, value interface{}) bool {
		if kernel, ok := value.(*Kernel); ok {
			kernel.Release()
		}
		return true
	})

	c.pool.Close()
	c.queue.Release()
	c.device.Release()
	c.adapter.Release()
	c.instance.Release()
	return nil
}

// Limits 返回设备限制（提供公开访问）
func (c *Context) Limits() wgpu.Limits {
	return c.limits
}

// Kernel 缓存的内核对象
type Kernel struct {
	Pipeline *wgpu.ComputePipeline
	Layout   *wgpu.BindGroupLayout
}

func (k *Kernel) Release() {
	if k.Pipeline != nil {
		k.Pipeline.Release()
	}
	if k.Layout != nil {
		k.Layout.Release()
	}
}
